from PyQt4 import QtCore , QtGui
from lib.papyDB import papyDB
from lib import papyDBTable
from lib.itemModel import treeItem
papyDB = papyDB()
    
class treeModel( QtCore.QAbstractItemModel ):    
    def __init__(self , headers ):
        super(treeModel, self).__init__()       
        self.rootItem = treeItem( 'root' )        
        self.setModelData( )
        
    def setModelData(self):  
        prjList = papyDBTable.PROJECT.search( 'prjCodeName' )      
        prjItems = [ treeItem( prj ) for prj in prjList ]        
        for prjCode , prjItem in zip( prjList , prjItems ):
            prjID = papyDBTable.PROJECT.search( 'prjID' , prjCodeName = prjCode )            
            seqList = papyDBTable.SEQUENCES.search( 'seqName' , prjID = prjID  )            
            seqItems = [ treeItem( seq ) for seq in seqList ]            
            prjItem.childItems.extend( seqItems )            
            for seqItem , seqName in zip( seqItems , seqList ) :
                seqItem.parentItem = prjItem                
                seqID = papyDBTable.SEQUENCES.search( 'seqID' , seqName = seqName , prjID = prjID )                
                shotList = papyDBTable.SHOT.search( 'shotName' , seqID = seqID )
                #print "@@@@", shotList
                if shotList != None:
                    shotItems = [ treeItem( shot )  for shot in shotList ]
                    seqItem.childItems.extend( shotItems )
                for shot in shotItems:
                    shot.parentItem = seqItem            
        self.rootItem.childItems.extend( prjItems ) 
        return len(prjList)
        
        
                        
    def columnCount(self, parent=QtCore.QModelIndex()):
        return self.rootItem.columnCount()        
    
    def data(self, index, role):
        if not index.isValid():
            return None
        if role != QtCore.Qt.DisplayRole and role != QtCore.Qt.EditRole:
            return None
        item = self.getItem(index)
        return item.data()
        
    def getItem(self, index):
        if index.isValid():
            item = index.internalPointer()
            if item:
                return item
        return self.rootItem
    
    def index(self, row, column, parent=QtCore.QModelIndex()):
        if parent.isValid() and parent.column() != 0:
            return QtCore.QModelIndex()
        parentItem = self.getItem(parent)
        childItem = parentItem.child(row)
        if childItem:
            return self.createIndex(row, column, childItem)
        else:
            return QtCore.QModelIndex()
    
    def parent(self, index):
        if not index.isValid():
            return QtCore.QModelIndex()
        childItem = self.getItem(index)
        parentItem = childItem.parent()
        if parentItem == self.rootItem:
            return QtCore.QModelIndex()
        elif parentItem == None :
            return QtCore.QModelIndex()
        return self.createIndex(parentItem.childNumber(), 0, parentItem )
    
    def rowCount(self, parent=QtCore.QModelIndex()):
        parentItem = self.getItem(parent)
        return parentItem.childCount()



class trWidget( QtGui.QWidget ):
    def __init__(self , parent = None):
        QtGui.QWidget.__init__(self, parent)
        self.treeItemModel = treeModel('test')
        
        self.resize(200,600)
        self.PSS_tree = QtGui.QTreeView(self)
        self.PSS_tree.setModel(self.treeItemModel)
        
        
#        for x in range(self.treeItemModel.setModelData()):
#            header_index = self.treeItemModel.index(x,0)
#            self.PSS_tree.setExpanded( header_index, True )
        self.PSS_tree.keyboardSearch('mini')
#        self.connect(self.PSS_tree.selectionModel(), QtCore.SIGNAL('selectionChanged(QItemSelection, QItemSelection)'), self.clickEvent)
        sm = self.PSS_tree.selectionModel()
        self.connect(sm, QtCore.SIGNAL('selectionChanged(QItemSelection, QItemSelection)'), self.selectionEvent)
#        print sm
        trLayout = QtGui.QGridLayout( self )
        trLayout.addWidget( self.PSS_tree )       
        self.show()

    def clickEvent(self):
        pass
#        print "CLICKED!"
        
        
    def selectionEvent(self, newSelection, oldSelection):
        pass
#        print "new=", newSelection.indexes()
        
#        self.params['test_category_id'] = int( self.model.itemFromIndex( newSelection.indexes()[0] ).text() ) 
#        self.params['test_category'] = self.model.itemFromIndex( newSelection.indexes()[2] ).text()
               
if __name__ == '__main__':     
    import sys    
    app = QtGui.QApplication(sys.argv)
    trWidget = trWidget()
    app.exec_()     

   